#include <iostream>
#include <string>
#include <cstring>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <net/if.h>
#include <sys/ioctl.h>
#include <linux/if_packet.h>
#include <net/ethernet.h>
#include <netinet/udp.h>
#include <netinet/ip.h>

#define DHCP_SERVER_PORT 67
#define DHCP_CLIENT_PORT 68
#define DHCP_MAGIC_COOKIE 0x63825363
#define DHCP_LEASE_TIME 86400 // 24 hours in seconds

#define DHCPDISCOVER 1
#define DHCPOFFER    2
#define DHCPREQUEST  3
#define DHCPACK      5

#define DHCP_OPTION_MSG_TYPE   53
#define DHCP_OPTION_SUBNET_MASK 1
#define DHCP_OPTION_ROUTER      3
#define DHCP_OPTION_DNS_SERVER  6
#define DHCP_OPTION_LEASE_TIME 51
#define DHCP_OPTION_END        255

struct dhcp_packet {
    uint8_t op;
    uint8_t htype;
    uint8_t hlen;
    uint8_t hops;
    uint32_t xid;
    uint16_t secs;
    uint16_t flags;
    uint32_t ciaddr;
    uint32_t yiaddr;
    uint32_t siaddr;
    uint32_t giaddr;
    uint8_t chaddr[16];
    uint8_t sname[64];
    uint8_t file[128];
    uint32_t magic_cookie;
    uint8_t options[308];
};

void add_dhcp_option(uint8_t *options, int *offset, uint8_t code, uint8_t len, const void *data) {
    options[(*offset)++] = code;
    options[(*offset)++] = len;
    memcpy(&options[*offset], data, len);
    *offset += len;
}

int main() {
    int sockfd;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);
    char buffer[1024];
    struct dhcp_packet *dhcp;

    if ((sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) < 0) {
        perror("socket");
        return 1;
    }

    int optval = 1;
    if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval)) < 0) {
        perror("setsockopt");
        close(sockfd);
        return 1;
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(DHCP_SERVER_PORT);
    server_addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("bind");
        close(sockfd);
        return 1;
    }

    std::cout << "DHCP Server started. Waiting for requests..." << std::endl;

    const char *assigned_ip = "192.168.1.2";
    const char *subnet_mask = "255.255.255.0";
    const char *router = "192.168.1.1";
    const char *dns_server = "8.8.8.8";

    while (1) {
        int bytes = recvfrom(sockfd, buffer, sizeof(buffer), 0, 
                           (struct sockaddr *)&client_addr, &client_len);
        if (bytes < 0) {
            perror("recvfrom");
            continue;
        }

        dhcp = (struct dhcp_packet *)buffer;

        if (ntohl(dhcp->magic_cookie) != DHCP_MAGIC_COOKIE) {
            continue;
        }

        uint8_t message_type = 0;
        uint8_t *options = dhcp->options;
        while (*options != DHCP_OPTION_END) {
            if (*options == DHCP_OPTION_MSG_TYPE) {
                message_type = options[2];
                break;
            }
            options += options[1] + 2;
        }

        switch (message_type) {
            case DHCPDISCOVER: {
                std::cout << "Received DHCPDISCOVER from client" << std::endl;

                memset(buffer, 0, sizeof(buffer));
                dhcp->op = BOOTREPLY;
                dhcp->yiaddr = inet_addr(assigned_ip);

                int offset = 0;
                uint8_t msg_type = DHCPOFFER;
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_MSG_TYPE, 1, &msg_type);
                
                uint32_t lease_time = htonl(DHCP_LEASE_TIME);
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_LEASE_TIME, 4, &lease_time);
                
                in_addr_t subnet = inet_addr(subnet_mask);
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_SUBNET_MASK, 4, &subnet);
                
                in_addr_t router_addr = inet_addr(router);
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_ROUTER, 4, &router_addr);
                
                in_addr_t dns_addr = inet_addr(dns_server);
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_DNS_SERVER, 4, &dns_addr);
                
                dhcp->options[offset++] = DHCP_OPTION_END;

                client_addr.sin_port = htons(DHCP_CLIENT_PORT);
                if (sendto(sockfd, buffer, sizeof(struct dhcp_packet), 0, 
                          (struct sockaddr *)&client_addr, client_len) < 0) {
                    perror("sendto");
                }
                break;
            }
            
            case DHCPREQUEST: {
                std::cout << "Received DHCPREQUEST from client" << std::endl;

                memset(buffer, 0, sizeof(buffer));
                dhcp->op = BOOTREPLY;
                dhcp->yiaddr = inet_addr(assigned_ip);

                int offset = 0;
                uint8_t msg_type = DHCPACK;
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_MSG_TYPE, 1, &msg_type);
                
                uint32_t lease_time = htonl(DHCP_LEASE_TIME);
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_LEASE_TIME, 4, &lease_time);
                
                in_addr_t subnet = inet_addr(subnet_mask);
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_SUBNET_MASK, 4, &subnet);
                
                in_addr_t router_addr = inet_addr(router);
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_ROUTER, 4, &router_addr);
                
                in_addr_t dns_addr = inet_addr(dns_server);
                add_dhcp_option(dhcp->options, &offset, DHCP_OPTION_DNS_SERVER, 4, &dns_addr);
                
                dhcp->options[offset++] = DHCP_OPTION_END;

                client_addr.sin_port = htons(DHCP_CLIENT_PORT);
                if (sendto(sockfd, buffer, sizeof(struct dhcp_packet), 0, 
                          (struct sockaddr *)&client_addr, client_len) < 0) {
                    perror("sendto");
                }
                
                std::cout << "Assigned IP address " << assigned_ip << " to client" << std::endl;
                break;
            }
            
            default:
                break;
        }
    }

    close(sockfd);
    return 0;
}